const std = @import("std");
const print = std.debug.print;
pub const libbpf = @cImport({
    @cInclude("libbpf.h");
    @cInclude("bpf.h");
    @cInclude("btf.h");
});
const config = @import("../config/config.zig");
const key = @import("../key/key.zig");

fn mergeAndPutSrcMap(src_map: *libbpf.struct_bpf_map, scidr: []const u8, id_tmp: u64) !void {
    var id = id_tmp;
    const sk = try key.parseCIDR(scidr);

    var prevId: u64 = 0;
    _ = libbpf.bpf_map__lookup_elem(src_map, &sk, @sizeOf(key.Ipv4LpmKey), &prevId, @sizeOf(@TypeOf(prevId)), 0);

    id |= prevId;

    const ret = libbpf.bpf_map__update_elem(src_map, &sk, @sizeOf(key.Ipv4LpmKey), &id, @sizeOf(u64), 0);
    if (ret != 0) {
        print("update ret: {d}\n", .{ret});
        return error.UpdateFailed;
    }

    print("Update src_map success\n", .{});
}

fn mergeAndPutDportMap(dport_map: *libbpf.struct_bpf_map, dport: u16, id_tmp: u64) !void {
    const dp = std.mem.nativeToBig(u16, dport);

    var id = id_tmp;

    var prevId: u64 = 0;
    _ = libbpf.bpf_map__lookup_elem(dport_map, &dp, @sizeOf(u16), &prevId, @sizeOf(@TypeOf(prevId)), 0);

    id |= prevId;

    const ret = libbpf.bpf_map__update_elem(dport_map, &dp, @sizeOf(u16), &id, @sizeOf(@TypeOf(id)), 0);
    if (ret != 0) {
        return error.UpdateFailed;
    }
    print("Update dport_map success\n", .{});
}
fn mergeAndPutProtoMap(proto_map: *libbpf.struct_bpf_map, proto: []const u8, id_tmp: u64) !void {
    var p: u8 = 0;
    if (std.mem.eql(u8, proto, "TCP")) {
        p = 6;
    } else if (std.mem.eql(u8, proto, "UDP")) {
        p = 17;
    } else {
        return error.UnsupportedProtocol;
    }

    var id = id_tmp;

    var prevId: u64 = 0;
    _ = libbpf.bpf_map__lookup_elem(proto_map, &p, @sizeOf(u8), &prevId, @sizeOf(@TypeOf(prevId)), 0);

    id |= prevId;
    const ret = libbpf.bpf_map__update_elem(proto_map, &p, @sizeOf(u8), &id, @sizeOf(@TypeOf(id)), 0);
    if (ret != 0) {
        return error.UpdateFailed;
    }

    print("Update proto_map success\n", .{});
}

fn putActionMap(action_map: *libbpf.struct_bpf_map, id: u64, action: []const u8) !void {
    var ac: u32 = 0;
    if (std.mem.eql(u8, action, "DROP")) {
        ac = 1;
    } else if (std.mem.eql(u8, action, "ACCEPT")) {
        ac = 2;
    } else {
        return error.UnsupportedAction;
    }

    const ret = libbpf.bpf_map__update_elem(action_map, &id, @sizeOf(u64), &ac, @sizeOf(@TypeOf(ac)), 0);
    if (ret != 0) {
        return error.UpdateFailed;
    }

    print("Update action_map success\n", .{});
}

fn addRule(obj: *libbpf.struct_bpf_object, rule: config.Rule, id: u64) !void {
    const src_map = libbpf.bpf_object__find_map_by_name(obj, "src_map").?;
    try mergeAndPutSrcMap(src_map, rule.scidr, id);
    const dport_map = libbpf.bpf_object__find_map_by_name(obj, "dport_map").?;
    try mergeAndPutDportMap(dport_map, rule.dport, id);
    const proto_map = libbpf.bpf_object__find_map_by_name(obj, "proto_map").?;
    try mergeAndPutProtoMap(proto_map, rule.proto, id);
    const action_map = libbpf.bpf_object__find_map_by_name(obj, "action_map").?;
    try putActionMap(action_map, id, rule.action);
}

pub fn processRules(obj: *libbpf.struct_bpf_object, rules: []const config.Rule) !void {
    for (rules, 0..) |rule, i| {
        const idx2 = @shlExact(@as(u64, 1), @intCast(i));
        print("Processing rule {d} {s} {d} {s} {s}\n", .{ idx2, rule.scidr, rule.dport, rule.proto, rule.action });
        addRule(obj, rule, idx2) catch |err| {
            print("Failed to add rule {d}: {any}\n", .{ idx2, err });
            return err;
        };
    }
}
